'''
    Author: Nebiyou Yismaw

    This is a python code that will do feature matching using OpenCV


'''

import numpy as np
import cv2

def main():
    img1 = cv2.imread("book.jpeg", cv2.IMREAD_GRAYSCALE)
    img2 = cv2.imread("book_scene.jpeg", cv2.IMREAD_GRAYSCALE)

    # Initiate ORB detector
    orb = cv2.ORB_create()

    # find the keypoints and descriptors with ORB
    kp1, des1 = orb.detectAndCompute(img1,None)
    kp2, des2 = orb.detectAndCompute(img2,None)

    # create BFMatcher object
    bf = cv2.BFMatcher(cv2.NORM_HAMMING, crossCheck=True)

    # Match descriptors.
    matches = bf.match(des1,des2)

    # Sort them in the order of their distance.
    matches = sorted(matches, key = lambda x:x.distance)

    # Draw first 10 matches.
    img3 = cv2.drawMatches(img1,kp1,img2,kp2,matches[:10],None,flags=cv2.DrawMatchesFlags_NOT_DRAW_SINGLE_POINTS)

    cv2.namedWindow('BFMatcher', cv2.WINDOW_NORMAL)
    cv2.imshow("BFMatcher",img3)
    cv2.resizeWindow('BFMatcher', 600,600)
    cv2.waitKey(0)

    # create Descriptor Matcher object
    matcher = cv2.DescriptorMatcher_create(cv2.DESCRIPTOR_MATCHER_BRUTEFORCE_HAMMING)
    # Match descriptors.
    matches = matcher.match(des1,des2)

    # Sort them in the order of their distance.
    matches = sorted(matches, key = lambda x:x.distance)

    # Draw first 10 matches.
    img3 = cv2.drawMatches(img1,kp1,img2,kp2,matches[:10],None)

    cv2.namedWindow('Descripter Matcher', cv2.WINDOW_NORMAL)
    cv2.imshow("Descripter Matcher", img3)
    cv2.resizeWindow('Descripter Matcher', 600,600)
    cv2.waitKey(0)

    # FLANN parameters
    FLANN_INDEX_KDTREE = 1
    index_params = dict(algorithm = FLANN_INDEX_KDTREE, trees = 5)
    search_params = dict(checks=50)   # or pass empty dictionary

    flann = cv2.FlannBasedMatcher(index_params,search_params)

    matches = flann.knnMatch(np.float32(des1),np.float32(des2),k=2)

    # Need to draw only good matches, so create a mask
    matchesMask = [[0,0] for i in range(len(matches))]

    # ratio test as per Lowe's paper
    for i,(m,n) in enumerate(matches):
        if m.distance < 0.7*n.distance:
            matchesMask[i]=[1,0]
    draw_params = dict(matchColor = (0,255,0),
                       singlePointColor = (255,0,0),
                       matchesMask = matchesMask,
                       flags = cv2.DrawMatchesFlags_DEFAULT)

    img3 = cv2.drawMatchesKnn(img1,kp1,img2,kp2,matches[:10],None)

    cv2.namedWindow('FLANN', cv2.WINDOW_NORMAL)
    cv2.imshow("FLANN",img3)
    cv2.resizeWindow('FLANN', 600,600)
    cv2.waitKey(0)

    cv2.destroyAllWindows()


if __name__=="__main__":
    main()
